package com.zz.inventory.service;

import com.alibaba.fastjson2.JSONObject;
import io.github.yuanbaobaoo.dify.client.DifyClientBuilder;
import io.github.yuanbaobaoo.dify.client.IDifyChatClient;
import io.github.yuanbaobaoo.dify.client.params.ParamMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;


@Service
@Slf4j
public class DIfyService {
    private final IDifyChatClient difyChatClient;

    public DIfyService(DifyClientBuilder difyClientBuilder) {
        this.difyChatClient = difyClientBuilder.buildChat();
    }

    public Flux<String> getStream(String userPrompt, String user) {
        Map<String, Object> inputs = new HashMap<>();
        inputs.put("test", "123");
        final ParamMessage paramMessage = ParamMessage.builder()
                .query(userPrompt)
                .user(user)
                .inputs(inputs)
                .build();

        return Flux.create(sink -> {
            final AtomicBoolean isCompleted = new AtomicBoolean(false);

            try {
                difyChatClient.sendMessagesAsync(paramMessage, chunk -> {
                    if (isCompleted.get()) return;

                    try {
                        // 1. 更安全的payload处理
                        JSONObject payload = chunk.getPayload();
                        if (payload == null) {
                            isCompleted.set(true);
                            sink.complete();
                            return;
                        }

                        Object answerObj = payload.get("answer");
                        if (answerObj == null) {
                            isCompleted.set(true);
                            sink.complete();
                            return;
                        }

                        String answer = answerObj.toString();
                        // 2. 确保格式正确
                        String message = answer.startsWith("data: ") ? answer : answer + "\n\n";
                        sink.next(message);

                        // 3. 结束标记检测
                        if (answer.contains("[DONE]")) {
                            isCompleted.set(true);
                            sink.complete();
                        }
                    } catch (Exception e) {
                        log.error("处理消息出错", e);
                        if (!isCompleted.get()) {
                            sink.next("data: [ERROR] 处理消息时出错: " + e.getMessage() + "\n\n");
                            isCompleted.set(true);
                            sink.complete();
                        }
                    }
                });
            } catch (Exception e) {
                log.error("API调用失败", e);
                sink.next("data: [ERROR] 服务调用失败: " + e.getMessage() + "\n\n");
                sink.complete();
            }

            sink.onDispose(() -> {
                isCompleted.set(true);
                log.info("客户端断开连接");
            });
        });
    }
}